import numpy as np
import matplotlib.pyplot as plt
from chaco.api import *
from traits.api import *
from traitsui.api import *
from enable.api import ComponentEditor
def water_er(f, T):
    ew0 = 87.134 - 0.19494 * T - 0.01276 * T ** 2 + 0.0002491 * T ** 3
    tw = 1.7681e-11 - 6.0861e-13 * T + 1.1042e-14 * T ** 2 - 8.1105e-17 * T ** 3
    ewomg = 4.9
    erw = (ewomg + (ew0 - ewomg) / (1 + (2 * np.pi * f * tw) ** 2)
           + 1j * 2 * np.pi * f * tw * (ew0 - ewomg) / (1 + (2 * np.pi * f * tw) ** 2))
    return erw

def seawater(f, T, SW):
    e0 = 8.854e-12
    ew0 = 87.134 - 0.19494 * T - 0.01276 * T ** 2 + 0.0002491 * T ** 3
    ewsea = ew0 * (1 + 1.613e-5 * T * SW - 3.656e-3 * SW + 3.210e-5 * SW ** 2 - 4.232e-7 * SW ** 3)
    tw = 1.7681e-11 - 6.0861e-13 * T + 1.1042e-14 * T ** 2 - 8.1105e-17 * T ** 3
    twsea = tw * (1 + 2.282e-5 * T * SW - 7.638e-4 * SW - 7.760e-6 * SW ** 2 + 1.105e-8 * SW ** 3)
    delta = 25 - T
    fai = delta * (2.033e-2 + 1.266e-4 * delta + 2.464e-6 * delta ** 2 - SW * (
                1.849e-5 - 2.551e-7 * delta + 2.551e-8 * delta ** 2))
    sigma = (1.8252e-1 * SW - 1.4619e-3 * SW ** 2 + 2.093e-5 * SW ** 3 - 1.282e-7 * SW ** 4) * np.exp(-fai)
    ewomg = 4.9
    erseaw = (ewomg + (ewsea - ewomg) / (1 + (2 * np.pi * f * twsea) ** 2)
              + 1j * 2 * np.pi * f * twsea * (ewsea - ewomg) / (1 + (2 * np.pi * f * twsea) ** 2)
              + 1j * sigma / (2 * np.pi * f * e0))
    return erseaw

def icewater(f, T):
    ewomg = 3.168
    ei0 = 203.168 + 2.5 * T + 0.15 * T**2
    sigma = 1.26 * np.exp(-0.6291 * 1e4 / (T + 273.15))
    alfa = 0.288 + 0.0052 * T + 0.00023 * T**2
    p = 9.990288e-5 * np.exp(0.6644 * 1e4 / (T + 273.15))
    e0 = 8.854e-12
    erice = (ewomg + (ei0 - ewomg) * (1 + p**(1-alfa) * np.sin(np.pi/2*alfa)) / (1 + 2*p**(1-alfa)*np.sin(np.pi/2*alfa) + p**(2*(1-alfa)))) + 1j * (p**(1-alfa) * np.cos(np.pi/2*alfa) * (ei0-ewomg) / (1+2*p**(1-alfa)*np.sin(np.pi/2*alfa) + p**(2*(1-alfa)))) + 1j * sigma / (2*np.pi*f*e0)
    return erice

def soil(mv, efw, fs, fc):
    rhoss = 2.65
    ess = 10.8143
    alfa = 0.65
    beta = 1.09 - 0.11 * fs + 0.18 * fc
    rhob = 2.04355 * (0.251 - 0.21 * fs + 0.22 * fc) ** 0.3018
    ersoil = (1 + rhob / rhoss * (ess ** alfa - 1) + mv ** beta * (efw ** alfa - 1)) ** (1 / alfa)
    return ersoil

def roughness(f, psai, dh):
    c=3.0e8
    erroughness = np.exp(-2*(2*np.pi*f*dh*np.sin(psai)/c)**2)*np.i0(2*(2*np.pi*f*dh*np.sin(psai)/c)**2)
    return erroughness

def vegetation(f, psai, k):
    c=3.0e8
    ervegetation = np.exp(-k*f*np.sin(psai)/c)*np.i0(k*f*np.sin(psai)/c)
    return ervegetation

class waterer(HasTraits):
    Temp = Range(low=0.0, high=70.0, value=20.0)
    data = Property(Array, observe=['Temp'])
    plot_left = Instance(Plot)
    plot_right = Instance(Plot)

    def _plot_left_default(self):
        f = np.arange(1, 101) * 1e9  # Hz
        T=self.Temp
        erw = water_er(f, T)
        er1=np.real(erw)
        plot = Plot(ArrayPlotData(T=f/1e9, E1=er1))
        plot.title = 'real permittivity of water'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "real permittivity"
        plot.plot(("T", "E1"), name='real permittivity', color=palette11[0])

        return plot
    def _plot_right_default(self):
        f = np.arange(1, 101) * 1e9  # Hz
        T=self.Temp
        erw = water_er(f, T)
        er2=np.imag(erw)
        plot = Plot(ArrayPlotData(T=f/1e9, E2=er2))
        plot.title = 'imag permittivity of water'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "imag permittivity"
        plot.plot(("T", "E2"), name='imag permittivity', color=palette11[0])
        return plot
    def _get_data(self):
        data={}
        f = np.arange(1, 101) * 1e9  # Hz
        T = self.Temp
        erw = water_er(f, T)
        er1 = np.real(erw)
        er2 = np.imag(erw)
        data['f']=f/1e9
        data['e1'] = er1
        data['e2'] = er2
        return data

    @observe("data")
    def _update_plot_left(self, event):
        self.plot_left.data.update_data(
            T=self.data['f'],
            E1=self.data['e1'],
        )
        self.plot_right.data.update_data(
            T=self.data['f'],
            E2=self.data['e2'],
        )
    traits_view = View(
        Group(
            HGroup(
                UItem(
                    "plot_left",
                    editor=ComponentEditor(),
                    resizable=True
                ),
                UItem(
                    "plot_right",
                    editor=ComponentEditor(),
                    resizable=True
                ),
            ),
        ),
        Item('Temp', label='水温度/℃', editor_args={'mode': 'slider'}),
    )
class seaer(HasTraits):
    Temp = Range(low=0.0, high=70.0, value=20.0)
    SW = Range(low=0.1, high=130.0, value=35.0)  # 盐度 千分比
    data = Property(Array, observe=['Temp','SW'])
    plot_left = Instance(Plot)
    plot_right = Instance(Plot)

    def _plot_left_default(self):
        f = np.arange(1, 101) * 1e9  # Hz
        T=self.Temp
        erseaw=seawater(f,T,self.SW)
        er1=np.real(erseaw)
        plot = Plot(ArrayPlotData(T=f/1e9, E1=er1))
        plot.title = 'real permittivity of sea water'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "real permittivity"
        plot.plot(("T", "E1"), name='real permittivity', color=palette11[0])

        return plot
    def _plot_right_default(self):
        f = np.arange(1, 101) * 1e9  # Hz
        T=self.Temp
        erseaw=seawater(f,T,self.SW)
        er2=np.imag(erseaw)
        plot = Plot(ArrayPlotData(T=f/1e9, E2=er2))
        plot.title = 'imag permittivity of sea water'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "imag permittivity"
        plot.plot(("T", "E2"), name='imag permittivity', color=palette11[0])
        return plot
    def _get_data(self):
        data={}
        f = np.arange(1, 101) * 1e9  # Hz
        T = self.Temp
        erseaw = seawater(f,T,self.SW)
        er1 = np.real(erseaw)
        er2 = np.imag(erseaw)
        data['f']=f/1e9
        data['e1'] = er1
        data['e2'] = er2
        return data

    @observe("data")
    def _update_plot_left(self, event):
        self.plot_left.data.update_data(
            T=self.data['f'],
            E1=self.data['e1'],
        )
        self.plot_right.data.update_data(
            T=self.data['f'],
            E2=self.data['e2'],
        )
    traits_view = View(
        Group(
            HGroup(
                UItem(
                    "plot_left",
                    editor=ComponentEditor(),
                    resizable=True
                ),
                UItem(
                    "plot_right",
                    editor=ComponentEditor(),
                    resizable=True
                ),
            ),
        ),
        Item('Temp', label='水温度/℃', editor_args={'mode': 'slider'}),
        Item('SW', label='盐度/千分比', editor_args={'mode': 'slider'}),
    )

class iceer(HasTraits):
    Temp = Range(low=-30, high=0.0, value=0.)
    data = Property(Array, observe=['Temp'])
    plot_left = Instance(Plot)
    plot_right = Instance(Plot)

    def _plot_left_default(self):
        f = np.arange(1, 101) * 1e9  # Hz
        T=self.Temp
        erice=icewater(f,T)
        er1=np.real(erice)
        plot = Plot(ArrayPlotData(T=f/1e9, E1=er1))
        plot.title = 'real permittivity of ice'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "real permittivity"
        plot.plot(("T", "E1"), name='real permittivity', color=palette11[0])

        return plot
    def _plot_right_default(self):
        f = np.arange(1, 101) * 1e9  # Hz
        T=self.Temp
        erice=icewater(f,T)
        er2=np.imag(erice)
        plot = Plot(ArrayPlotData(T=f/1e9, E2=er2))
        plot.title = 'imag permittivity of ice'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "imag permittivity"
        plot.plot(("T", "E2"), name='imag permittivity', color=palette11[0])
        return plot
    def _get_data(self):
        data={}
        f = np.arange(1, 101) * 1e9  # Hz
        T = self.Temp
        erice=icewater(f,T)
        er1 = np.real(erice)
        er2 = np.imag(erice)
        data['f']=f/1e9
        data['e1'] = er1
        data['e2'] = er2
        return data

    @observe("data")
    def _update_plot_left(self, event):
        self.plot_left.data.update_data(
            T=self.data['f'],
            E1=self.data['e1'],
        )
        self.plot_right.data.update_data(
            T=self.data['f'],
            E2=self.data['e2'],
        )
    traits_view = View(
        Group(
            HGroup(
                UItem(
                    "plot_left",
                    editor=ComponentEditor(),
                    resizable=True
                ),
                UItem(
                    "plot_right",
                    editor=ComponentEditor(),
                    resizable=True
                ),
            ),
        ),
        Item('Temp', label='水温度/℃', editor_args={'mode': 'slider'}),
    )

class soiler(HasTraits):
    Temp = Range(low=0.0, high=70.0, value=20.0)
    mv = Range(low=0.0, high=1.0, value=0.5)  # 土壤湿度
    fs = Range(low=0.0, high=1.0, value=0.8)  # 沙质量比
    fc = Range(low=0.0, high=1.0, value=0.2)  # 黏土质量比
    data = Property(Array, observe=['Temp','mv','fs','fc'])
    plot_left = Instance(Plot)
    plot_right = Instance(Plot)

    def _plot_left_default(self):
        f = np.arange(1, 40.1,0.1) * 1e9  # Hz
        T=self.Temp
        efw = water_er(f, T)
        ersoil = soil(self.mv, efw, self.fs, self.fc)
        er1=np.real(ersoil)
        plot = Plot(ArrayPlotData(T=f/1e9, E1=er1))
        plot.title = 'real permittivity of soil'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "real permittivity"
        plot.plot(("T", "E1"), name='real permittivity', color=palette11[0])

        return plot
    def _plot_right_default(self):
        f = np.arange(1, 40.1,0.1) * 1e9  # Hz
        T=self.Temp
        efw = water_er(f, T)
        ersoil = soil(self.mv, efw, self.fs, self.fc)
        er2=np.imag(ersoil)
        plot = Plot(ArrayPlotData(T=f/1e9, E2=er2))
        plot.title = 'imag permittivity of soil'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "imag permittivity"
        plot.plot(("T", "E2"), name='imag permittivity', color=palette11[0])
        return plot
    def _get_data(self):
        data={}
        f = np.arange(1, 40.1,0.1) * 1e9  # Hz
        T = self.Temp
        efw = water_er(f, T)
        ersoil=soil(self.mv,efw,self.fs,self.fc)
        er1 = np.real(ersoil)
        er2 = np.imag(ersoil)
        data['f']=f/1e9
        data['e1'] = er1
        data['e2'] = er2
        return data

    @observe("data")
    def _update_plot_left(self, event):
        self.plot_left.data.update_data(
            T=self.data['f'],
            E1=self.data['e1'],
        )
        self.plot_right.data.update_data(
            T=self.data['f'],
            E2=self.data['e2'],
        )
    traits_view = View(
        Group(
            HGroup(
                UItem(
                    "plot_left",
                    editor=ComponentEditor(),
                    resizable=True
                ),
                UItem(
                    "plot_right",
                    editor=ComponentEditor(),
                    resizable=True
                ),
            ),
        ),
        Item('Temp', label='土壤温度/℃', editor_args={'mode': 'slider'}),
        Item('mv', label='土壤湿度', editor_args={'mode': 'slider'}),
        Item('fs', label='沙质量比', editor_args={'mode': 'slider'}),
        Item('fc', label='黏土质量比', editor_args={'mode': 'slider'}),
    )

class roughnesser(HasTraits):
    psai = Range(low=0.0, high=90.0, value=20.0)
    dh = Range(low=0.0, high=1.0, value=0.01)  # 表面粗糙度的均方根高度差
    data = Property(Array, observe=['psai','dh'])
    plot_left = Instance(Plot)

    def _plot_left_default(self):
        f = np.arange(1, 40.1,0.1) * 1e9  # Hz
        erroughness = roughness(f, self.psai, self.dh)
        plot = Plot(ArrayPlotData(T=f/1e9, E1=erroughness))
        plot.title = 'Surface roughness coefficient'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "Surface roughness coefficient"
        plot.plot(("T", "E1"), name='地表粗糙度系数', color=palette11[0])

        return plot

    def _get_data(self):
        data={}
        f = np.arange(1, 40.1,0.1) * 1e9  # Hz
        erroughness = roughness(f, self.psai, self.dh)
        data['f']=f/1e9
        data['e1'] = erroughness
        return data

    @observe("data")
    def _update_plot_left(self, event):
        self.plot_left.data.update_data(
            T=self.data['f'],
            E1=self.data['e1'],
        )

    traits_view = View(
        Group(
            HGroup(
                UItem(
                    "plot_left",
                    editor=ComponentEditor(),
                    resizable=True
                ),

            ),
        ),
        Item('psai', label='雷达波掠射角', editor_args={'mode': 'slider'}),
        Item('dh', label='表面粗糙度的均方根高度差', editor_args={'mode': 'slider'}),
    )

class vegetationer(HasTraits):
    psai = Range(low=0.0, high=90.0, value=20.0)
    k = Range(low=0.0, high=10.0, value=1)  # 植被覆盖程度
    data = Property(Array, observe=['psai','k'])
    plot_left = Instance(Plot)

    def _plot_left_default(self):
        f = np.arange(1, 40.1,0.1) * 1e9  # Hz
        ervegetation = vegetation(f, self.psai, self.k)
        plot = Plot(ArrayPlotData(T=f/1e9, E1=ervegetation))
        plot.title = 'Vegetation reflection coefficient'
        plot.x_axis.title = "Frequency /GHz"
        plot.y_axis.title = "Vegetation reflection coefficient"
        plot.plot(("T", "E1"), name='植被反射系数', color=palette11[0])

        return plot

    def _get_data(self):
        data={}
        f = np.arange(1, 40.1,0.1) * 1e9  # Hz
        ervegetation = vegetation(f, self.psai, self.k)
        data['f']=f/1e9
        data['e1'] = ervegetation
        return data

    @observe("data")
    def _update_plot_left(self, event):
        self.plot_left.data.update_data(
            T=self.data['f'],
            E1=self.data['e1'],
        )

    traits_view = View(
        Group(
            HGroup(
                UItem(
                    "plot_left",
                    editor=ComponentEditor(),
                    resizable=True
                ),

            ),
        ),
        Item('psai', label='雷达波掠射角', editor_args={'mode': 'slider'}),
        Item('k', label='植被覆盖程度', editor_args={'mode': 'slider'}),
    )

class Form(HasTraits):
    waterer = Instance(waterer, ())
    seaer = Instance(seaer, ())
    iceer = Instance(iceer, ())
    soiler = Instance(soiler, ())
    roughnesser = Instance(roughnesser, ())
    vegetationer = Instance(vegetationer, ())
    traits_view = View(
        Tabbed(
            UItem('waterer', style='custom', label='水的介电常数'),
            UItem('seaer', style='custom', label='海水的介电常数'),
            UItem('iceer', style='custom', label='冰的介电常数'),
            UItem('soiler', style='custom', label='土壤的介电常数'),
            UItem('roughnesser', style='custom', label='地表粗糙度系数'),
            UItem('vegetationer', style='custom', label='植被反射系数'),
        ),
        title='地物电磁参数计算',
        width=1200,
        height=800,
        resizable=True,
    )

if __name__ == '__main__':
    Form().configure_traits()